#!/usr/bin/env python3
"""
Inspect DREAMER.mat structure (supports MATLAB v7 and v7.3/HDF5) and output a Markdown summary.

Usage:
  python inspect_dreamer_mat.py /path/to/DREAMER.mat [--output output.md] [--max-items 20]

This script attempts scipy.io.loadmat first; if the file is HDF5 (v7.3), it falls back to h5py.
The output Markdown includes dataset paths, shapes, dtypes and truncated previews.
"""

import argparse
import json
import sys
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple


def human_shape(shape: Iterable[int]) -> str:
    try:
        return "x".join(str(int(s)) for s in shape)
    except Exception:
        return str(shape)


def truncate_list(values: Iterable[Any], max_items: int) -> str:
    items: List[str] = []
    count = 0
    for v in values:
        count += 1
        if count > max_items:
            break
        try:
            if isinstance(v, (bytes, bytearray)):
                items.append(repr(v[:16]))
            else:
                items.append(repr(v))
        except Exception:
            items.append("<unrepr>")
    suffix = " …" if count > max_items else ""
    return ", ".join(items) + suffix


def inspect_scipy(mat_path: Path, max_items: int) -> Tuple[str, Dict[str, Any]]:
    from scipy.io import loadmat, matlab

    data = loadmat(str(mat_path), struct_as_record=False, squeeze_me=True)
    # Remove MATLAB loader metadata keys
    for meta in ("__header__", "__version__", "__globals__"):
        data.pop(meta, None)

    lines: List[str] = []
    lines.append(f"# DREAMER.mat 结构 (scipy.io.loadmat)\n")

    def describe(name: str, obj: Any, indent: int = 0) -> None:
        pad = "  " * indent
        if hasattr(obj, "dtype") and hasattr(obj, "shape"):
            try:
                dt = str(obj.dtype)
                sh = human_shape(getattr(obj, "shape", ()))
                preview = None
                try:
                    import numpy as np
                    flat = obj.ravel()
                    preview = truncate_list(flat.tolist(), max_items)
                except Exception:
                    preview = "<unavailable>"
                lines.append(f"{pad}- {name}: ndarray shape={sh}, dtype={dt}, preview=[{preview}]")
            except Exception:
                lines.append(f"{pad}- {name}: ndarray <uninspectable>")
            return

        if isinstance(obj, (list, tuple)):
            lines.append(f"{pad}- {name}: {type(obj).__name__} len={len(obj)}")
            for idx, item in enumerate(obj[: max_items]):
                describe(f"[{idx}]", item, indent + 1)
            if len(obj) > max_items:
                lines.append(f"{pad}  … ({len(obj) - max_items} more)")
            return

        if isinstance(obj, dict):
            lines.append(f"{pad}- {name}: dict keys={len(obj)}")
            for k, v in list(obj.items())[: max_items]:
                describe(f"{k}", v, indent + 1)
            if len(obj) > max_items:
                lines.append(f"{pad}  … ({len(obj) - max_items} more)")
            return

        if isinstance(obj, matlab.mio5_params.mat_struct):
            # MATLAB struct loaded as an object with attributes
            fields = [f for f in dir(obj) if not f.startswith("_")]
            lines.append(f"{pad}- {name}: mat_struct fields={len(fields)}")
            shown = 0
            for f in fields:
                if shown >= max_items:
                    break
                try:
                    val = getattr(obj, f)
                except Exception:
                    val = "<unavailable>"
                describe(f, val, indent + 1)
                shown += 1
            if len(fields) > max_items:
                lines.append(f"{pad}  … ({len(fields) - max_items} more)")
            return

        # Fallback for scalars/strings
        try:
            lines.append(f"{pad}- {name}: {type(obj).__name__} value={repr(obj)[:200]}")
        except Exception:
            lines.append(f"{pad}- {name}: {type(obj).__name__} <unrepr>")

    for key, value in data.items():
        describe(key, value, 0)

    return ("scipy", {"lines": lines})


def inspect_h5(mat_path: Path, max_items: int) -> Tuple[str, Dict[str, Any]]:
    import h5py

    lines: List[str] = []
    lines.append(f"# DREAMER.mat 结构 (HDF5 v7.3 via h5py)\n")

    def attrs_to_str(attrs: h5py.AttributeManager) -> str:
        parts: List[str] = []
        for k in attrs.keys():
            try:
                v = attrs[k]
                # Uniformize bytes
                if isinstance(v, (bytes, bytearray)):
                    v = v.decode(errors="ignore")
                parts.append(f"{k}={repr(v)[:80]}")
            except Exception:
                parts.append(f"{k}=<unavailable>")
        return ", ".join(parts)

    def visit(name: str, obj: h5py.Dataset | h5py.Group) -> None:
        if isinstance(obj, h5py.Dataset):
            shape = human_shape(obj.shape)
            dtype = str(obj.dtype)
            preview = ""
            try:
                # Read a tiny slice safely
                if obj.ndim == 0:
                    sample = obj[()]
                else:
                    # Build a slice that limits first axis to max_items
                    index = [slice(0, min(max_items, obj.shape[0]))] + [slice(0, 1) for _ in range(max(0, obj.ndim - 1))]
                    sample = obj[tuple(index)]
                # Convert to JSON-like preview
                try:
                    import numpy as np
                    preview = truncate_list(np.array(sample).ravel().tolist(), max_items)
                except Exception:
                    preview = "<unavailable>"
            except Exception:
                preview = "<no access>"
            lines.append(f"- /{name}: dataset shape={shape}, dtype={dtype}, attrs=[{attrs_to_str(obj.attrs)}], preview=[{preview}]")
        else:
            lines.append(f"- /{name}: group attrs=[{attrs_to_str(obj.attrs)}]")

    with h5py.File(str(mat_path), "r") as f:
        # Root
        lines.append(f"- /: file attrs=[{attrs_to_str(f.attrs)}]")
        f.visititems(lambda n, o: visit(n, o))

    return ("h5py", {"lines": lines})


def main() -> int:
    parser = argparse.ArgumentParser(description="Inspect DREAMER.mat and output Markdown structure")
    parser.add_argument("mat_file", type=str, help="Path to DREAMER.mat")
    parser.add_argument("--output", type=str, default=None, help="Output Markdown path (default: alongside .mat)")
    parser.add_argument("--max-items", type=int, default=20, help="Max items per collection/preview")
    args = parser.parse_args()

    mat_path = Path(args.mat_file).expanduser().resolve()
    if not mat_path.exists():
        print(f"File not found: {mat_path}", file=sys.stderr)
        return 1

    output_md = Path(args.output) if args.output else mat_path.with_name(mat_path.stem + "_structure.md")

    inspectors = []  # type: List[Tuple[str, Dict[str, Any]]]
    # Try scipy first
    try:
        backend, payload = inspect_scipy(mat_path, args.max_items)
        inspectors.append((backend, payload))
    except Exception as e:
        inspectors.append(("scipy", {"error": repr(e)}))
        # Try h5py fallback
        try:
            backend, payload = inspect_h5(mat_path, args.max_items)
            inspectors.append((backend, payload))
        except Exception as e2:
            inspectors.append(("h5py", {"error": repr(e2)}))

    # Compose Markdown
    md_lines: List[str] = []
    md_lines.append(f"# DREAMER.mat 数据结构概览\n")
    md_lines.append(f"- **文件**: `{mat_path}`\n")
    md_lines.append(f"- **生成时间**: automatic\n")
    md_lines.append("")

    for backend, payload in inspectors:
        md_lines.append(f"## 后端: {backend}\n")
        if "error" in payload:
            md_lines.append(f"> 解析失败: `{payload['error']}`\n")
        else:
            md_lines.extend(payload["lines"])
        md_lines.append("")

    output_md.parent.mkdir(parents=True, exist_ok=True)
    output_md.write_text("\n".join(md_lines), encoding="utf-8")
    print(f"Wrote: {output_md}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())


